﻿#ifndef __SS_RB_TREE_H__
#define __SS_RB_TREE_H__
#include "ssRBTreeNode.h"
#include "../../ssAssert.h"
#include "../../ssCompareSize.h"
#include "../../memory/ssDefaultMalloc.h"
#include "../ssStack.h"
#include "../ssPair.h"

#define SS_RBT_OK					0
#define SS_RBT_FAILED				1
#define SS_RBT_CREATE_NODE_FAILED	2	//创建新节点失败
#define SS_RBT_KEY_EXISTED			3	//键值已存在
#define SS_RBT_KEY_NOT_FOUND		4	//键值不存在
#define SS_RBT_DEL_NULL				5	//删除空节点
#define SS_RBT_DEL_ERROR			6	//删除节点出错
#define SS_RBT_DEL_ILLEGAL			7	//非法删除节点

namespace StarSeeker
{

template<typename T_KEY, typename T_POOL=ssDefMallocObject<ssRBTreeNode<T_KEY> > >
class ssRBTree
{
public:
	ssRBTree(void);
	ssRBTree(const ssRBTree& t);
	virtual ~ssRBTree(void);
	void operator= (const ssRBTree& t);

	//插入
	ssPair<ssRBTreeNode<T_KEY>*, error_t> Insert(const T_KEY& k);
	//查找
	ssRBTreeNode<T_KEY>* Find(const T_KEY& k);
	//删除
	error_t Remove(ssRBTreeNode<T_KEY>* node);
	error_t Remove(const T_KEY& k);
	//最小值
	ssRBTreeNode<T_KEY>* GetMin();
	//最大值
	ssRBTreeNode<T_KEY>* GetMax();
	//返回节点数
	ulint GetCount() {return m_count;}
	//清空树
	void Empty();
	//错误检查
	bool Check();
	//遍历
	ssRBTreeNode<T_KEY>* Begin();
	ssRBTreeNode<T_KEY>* End();
	ssRBTreeNode<T_KEY>* Next();

private:
	//复制对象
	void CopyFrom(const ssRBTree& t);
	//创建节点
	ssRBTreeNode<T_KEY>* CreateNode();
	//删除节点
	void DeleteNode(ssRBTreeNode<T_KEY>* itr);
	//左旋
	bool LeftRotate(ssRBTreeNode<T_KEY>* x);
	//右旋
	bool RightRotate(ssRBTreeNode<T_KEY>* y);
	//插入修复
	void InsertFixUp(ssRBTreeNode<T_KEY>* itr);
	//删除时找到替换点
	ssRBTreeNode<T_KEY>* FindSuccessor(ssRBTreeNode<T_KEY>* itr);
	//删除修复
	void DeleteFixUp(ssRBTreeNode<T_KEY>* node, ssRBTreeNode<T_KEY>* parent);
    //检查红色节点
    bool CheckRed(ssRBTreeNode<T_KEY>* itr);
	//检查黑色节点
	bool CheckBlack(ssRBTreeNode<T_KEY>* itr, bool root);
	//遍历删除节点
	void RecursionEmpty(ssRBTreeNode<T_KEY>* itr);
	//遍历复制
	void RecursionCopy(const ssRBTreeNode<T_KEY>* itr);

private:
	ssRBTreeNode<T_KEY>*				m_root;
	ulint								m_count;
	T_POOL								m_objPool;
	ssNumStack<ssRBTreeNode<T_KEY>*>	m_stack;
	ssRBTreeNode<T_KEY>*				m_iterator;
	ulint								m_checkBlackCount;
};


template<typename T_KEY, typename T_POOL>
inline ssRBTree<T_KEY, T_POOL>::ssRBTree(void)
{
	m_root = 0;
	m_count = 0;
	m_iterator = 0;
}

template<typename T_KEY, typename T_POOL>
inline ssRBTree<T_KEY, T_POOL>::ssRBTree(const ssRBTree & t)
{
	m_root = 0;
	//m_null.m_color = SS_RBT_BLACK;
	m_count = 0;
	m_iterator = 0;
	CopyFrom(t);
}

template<typename T_KEY, typename T_POOL>
inline ssRBTree<T_KEY, T_POOL>::~ssRBTree(void)
{
	Empty();
}

template<typename T_KEY, typename T_POOL>
void StarSeeker::ssRBTree<T_KEY, T_POOL>::operator=(const ssRBTree& t)
{
	Empty();
	CopyFrom(t);
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::CreateNode()
{
	return (ssRBTreeNode<T_KEY>*)m_objPool.Malloc();
}

template<typename T_KEY,typename T_POOL>
inline void ssRBTree<T_KEY,T_POOL>::DeleteNode(ssRBTreeNode<T_KEY>* itr)
{
	itr->m_pOnwer = 0;
	m_objPool.Free(itr);
}

template<typename T_KEY,typename T_POOL>
inline bool ssRBTree<T_KEY,T_POOL>::LeftRotate(ssRBTreeNode<T_KEY>* x)
{
	if(x != 0)
	{
		ssRBTreeNode<T_KEY>* y = x->m_right;
		if(y != 0)
		{
			//y的左子成为x的右子
			x->m_right = y->m_left;
			//y成为x的父结点
			if (y->m_left != 0)
			{
				y->m_left->m_parent = x;
			}
			y->m_parent = x->m_parent;
			//如果是根节点
			if(x->m_parent == 0)
			{
				m_root = y;
			}
			else
			{
				if (x->m_parent->m_left==x) //将y挂到x父节点的左子或右子
				{
					x->m_parent->m_left = y;
				}
				else if(x->m_parent->m_right==x)
				{
					x->m_parent->m_right = y;
				}
				else
				{
					SS_ASSERT(0);
				}
			}
			x->m_parent = y;
			y->m_left = x;
			return true;
		}
	}
	return false;
}

template<typename T_KEY,typename T_POOL>
inline bool ssRBTree<T_KEY,T_POOL>::RightRotate(ssRBTreeNode<T_KEY>* y)
{
	if(y != 0)
	{
		ssRBTreeNode<T_KEY>* x = y->m_left;
		if(x != 0)
		{
			//y的右子成为x的左子
			y->m_left = x->m_right;
			//y成为x的父结点
			if (x->m_right != 0)
			{
				x->m_right->m_parent = y;
			}
			x->m_parent = y->m_parent;
			//如果是根节点
			if(y->m_parent == 0)
			{
				m_root = x;
			}
			else
			{
				if (y->m_parent->m_left==y) //将y挂到x父节点的左子或右子
				{
					y->m_parent->m_left = x;
				}
				else if(y->m_parent->m_right==y)
				{
					y->m_parent->m_right = x;
				}
				else
				{
					SS_ASSERT(0);
				}
			}
			y->m_parent = x;
			x->m_right = y;
			return true;
		}
	}
	return false;
}

template<typename T_KEY,typename T_POOL>
inline ssPair<ssRBTreeNode<T_KEY>*, error_t> ssRBTree<T_KEY,T_POOL>::Insert(const T_KEY& k)
{
	ssPair<ssRBTreeNode<T_KEY>*, error_t> ret(0,1);

	//查找插入点
	ssRBTreeNode<T_KEY>* index = m_root;
	ssRBTreeNode<T_KEY>* insert = 0;
	while (index != 0)
	{
		insert = index;
		if(ssCompLess(k, index->m_key))
		{
			index = index->m_left;
		}
		else if(ssCompGreater(k, index->m_key))
		{
			index = index->m_right;
		}
		else //节点已存在
		{
			ret.first = insert;
			ret.second = SS_RBT_KEY_EXISTED;
			return ret;
		}
	}
	//新节点
	ssRBTreeNode<T_KEY>* newNode = CreateNode();
	if(!newNode)
	{
		ret.second = SS_RBT_CREATE_NODE_FAILED;
		return ret;
	}
	newNode->m_color = SS_RBT_RED; //初始红色
	newNode->m_key = k;
	newNode->m_right = 0;
	newNode->m_left = 0;
	newNode->m_pOnwer = this;
	if (insert==0) //如果是空树
	{
		m_root = newNode;
		m_root->m_color = SS_RBT_BLACK;
		m_root->m_parent = 0;
	}
	else
	{
		if(k<insert->m_key)
		{
			insert->m_left = newNode;
		}
		else if(k>insert->m_key)
		{
			insert->m_right = newNode;
		}
		else
		{
			SS_ASSERT(0);
		}
		newNode->m_parent = insert;
	}
	InsertFixUp(newNode);
	++m_count;

	ret.first = newNode;
	ret.second = SS_RBT_OK;
	return ret;
}

template<typename T_KEY,typename T_POOL>
inline void ssRBTree<T_KEY,T_POOL>::InsertFixUp(ssRBTreeNode<T_KEY>* itr)
{
	while (itr->m_parent && itr->m_parent->m_color==SS_RBT_RED) //父节点红色才需要调整
	{
		if (itr->m_parent == itr->m_parent->m_parent->m_left) //父节点是祖父的左子
		{
			ssRBTreeNode<T_KEY>* uncle = itr->m_parent->m_parent->m_right; //祖父的右子就是叔叔
			if (uncle && uncle->m_color == SS_RBT_RED) //插入情况1，z的叔叔y是红色的
			{
				itr->m_parent->m_color = SS_RBT_BLACK;
				uncle->m_color = SS_RBT_BLACK;
				itr->m_parent->m_parent->m_color = SS_RBT_RED;
				itr = itr->m_parent->m_parent;
				continue;
			}
			else//插入情况2：z的叔叔y是黑色的
			{
				if (itr == itr->m_parent->m_right) //且z是右子
				{
					itr = itr->m_parent;
					LeftRotate(itr);
				}
				if (itr == itr->m_parent->m_left) //插入情况3：z的叔叔y是黑色的，但z是左子
				{
					itr->m_parent->m_color = SS_RBT_BLACK;
					itr->m_parent->m_parent->m_color = SS_RBT_RED;
					RightRotate(itr->m_parent->m_parent);
				}
			}
		}
		else if (itr->m_parent == itr->m_parent->m_parent->m_right)//父节点是祖父的右子
		{
			ssRBTreeNode<T_KEY>* uncle = itr->m_parent->m_parent->m_left;
			if (uncle && uncle->m_color == SS_RBT_RED)
			{
				itr->m_parent->m_color = SS_RBT_BLACK;
				uncle->m_color = SS_RBT_BLACK;
				itr->m_parent->m_parent->m_color = SS_RBT_RED;
				itr = itr->m_parent->m_parent;
			}
			else
			{
				if (itr == itr->m_parent->m_left)
				{
					itr = itr->m_parent;
					RightRotate(itr);     //与上述代码相比，左旋改为右旋
				}
				if (itr == itr->m_parent->m_right)
				{
					itr->m_parent->m_color = SS_RBT_BLACK;
					itr->m_parent->m_parent->m_color = SS_RBT_RED;
					LeftRotate(itr->m_parent->m_parent);   //右旋改为左旋，即可。
				}
			}
		}
		else
		{
			SS_ASSERT(0);
		}
	}
	m_root->m_color = SS_RBT_BLACK;
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::Find(const T_KEY& k)
{
	ssRBTreeNode<T_KEY>* index = m_root;
	while(index != 0)
	{
		if(ssCompLess(k, index->m_key))
		{
			index  = index->m_left;  //比当前的小，往左
		}
		else if(ssCompGreater(k, index->m_key))
		{
			index = index->m_right;  //比当前的大，往右
		}
		else
		{
			break;
		}
	}
	return index;
}

template<typename T_KEY,typename T_POOL>
inline error_t ssRBTree<T_KEY,T_POOL>::Remove(ssRBTreeNode<T_KEY>* node)
{
	if(!node || node ==0)
	{
		return SS_RBT_DEL_NULL;
	}
	if (node->m_pOnwer!=this)
	{
		return SS_RBT_DEL_ILLEGAL;
	}

	ssRBTreeNode<T_KEY> *child, *parent;
	bool color;

	// 被删除节点的"左右孩子都不为空"的情况。
	if ((node->m_left != NULL) && (node->m_right != NULL))
	{
		// 被删节点的后继节点。(称为"取代节点")
		// 用它来取代"被删节点"的位置，然后再将"被删节点"去掉。
		ssRBTreeNode<T_KEY> *replace = node;

		// 获取后继节点
		replace = replace->m_right;
		while (replace->m_left != NULL)
			replace = replace->m_left;

		// "node节点"不是根节点(只有根节点不存在父节点)
		if (node->m_parent)
		{
			if (node->m_parent->m_left == node)
				node->m_parent->m_left = replace;
			else
				node->m_parent->m_right = replace;
		}
		else
			// "node节点"是根节点，更新根节点。
			m_root = replace;

		// child是"取代节点"的右孩子，也是需要"调整的节点"。
		// "取代节点"肯定不存在左孩子！因为它是一个后继节点。
		child = replace->m_right;
		parent = replace->m_parent;
		// 保存"取代节点"的颜色
		color = replace->m_color;

		// "被删除节点"是"它的后继节点的父节点"
		if (parent == node)
		{
			parent = replace;
		}
		else
		{
			// child不为空
			if (child)
				child->m_parent = parent;
			parent->m_left = child;

			replace->m_right = node->m_right;
			node->m_right->m_parent = replace;
		}

		replace->m_parent = node->m_parent;
		replace->m_color = node->m_color;
		replace->m_left = node->m_left;
		node->m_left->m_parent = replace;

		if (color == SS_RBT_BLACK)
			DeleteFixUp(child, parent);

		DeleteNode(node);
		--m_count;
		return SS_RBT_OK;
	}

	if (node->m_left != NULL)
		child = node->m_left;
	else
		child = node->m_right;

	parent = node->m_parent;
	// 保存"取代节点"的颜色
	color = node->m_color;

	if (child)
		child->m_parent = parent;

	// "node节点"不是根节点
	if (parent)
	{
		if (parent->m_left == node)
			parent->m_left = child;
		else
			parent->m_right = child;
	}
	else
		m_root = child;

	if (color == SS_RBT_BLACK)
		DeleteFixUp(child, parent);
	
	DeleteNode(node);
	--m_count;
	return SS_RBT_OK;
}

template<typename T_KEY,typename T_POOL>
inline error_t ssRBTree<T_KEY,T_POOL>::Remove(const T_KEY& k)
{
	ssRBTreeNode<T_KEY>* pDelete = Find(k);
	return Remove(pDelete);
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::FindSuccessor(ssRBTreeNode<T_KEY>* itr)
{
	if(itr==0 || itr->m_left==0 || itr->m_right==0)
		return 0;

	ssRBTreeNode<T_KEY>* result = itr->m_left; //获取左子
	while(result != 0) //找出左子树中右值的最大值
	{
		if(result->m_right != 0)
		{
			result = result->m_right;
		}
		else
		{
			break;
		}
	}
	//如果左子树的右值为空，查看右子树
	if (result == itr->m_left)
	{
		result = itr->m_right; //获取右子
		while(result != 0) //找出右子树中左值的最小值
		{
			if(result->m_left != 0)
			{
				result = result->m_left;
			}
			else
			{
				break;
			}
		}
		if (result == itr->m_right)
		{
			result = itr->m_left;
		}
	}
	return result;
}

template<typename T_KEY,typename T_POOL>
inline void ssRBTree<T_KEY,T_POOL>::DeleteFixUp(ssRBTreeNode<T_KEY>* node, ssRBTreeNode<T_KEY>* parent)
{
	ssRBTreeNode<T_KEY>* other = 0;

	while ((!node || node->m_color== SS_RBT_BLACK) && node != m_root)
	{
		if (parent->m_left == node)
		{
			other = parent->m_right;
			if (other->m_color == SS_RBT_RED)
			{
				// Case 1: x的兄弟w是红色的  
				other->m_color = SS_RBT_BLACK;
				parent->m_color = SS_RBT_RED;
				LeftRotate(parent);
				other = parent->m_right;
			}
			if ((!other->m_left || other->m_left->m_color == SS_RBT_BLACK) &&
				(!other->m_right || other->m_right->m_color == SS_RBT_BLACK))
			{
				// Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的  
				other->m_color = SS_RBT_RED;
				node = parent;
				parent = node->m_parent;
			}
			else
			{
				if (!other->m_right || other->m_right->m_color == SS_RBT_BLACK)
				{
					// Case 3: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。  
					other->m_left->m_color = SS_RBT_BLACK;
					other->m_color = SS_RBT_RED;
					RightRotate(other);
					other = parent->m_right;
				}
				// Case 4: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。
				other->m_color = parent->m_color;
				parent->m_color = SS_RBT_BLACK;
				other->m_right->m_color = SS_RBT_BLACK;
				LeftRotate(parent);
				node = m_root;
				break;
			}
		}
		else
		{
			other = parent->m_left;
			if (other->m_color == SS_RBT_RED)
			{
				// Case 1: x的兄弟w是红色的  
				other->m_color = SS_RBT_BLACK;
				parent->m_color = SS_RBT_RED;
				RightRotate(parent);
				other = parent->m_left;
			}
			if ((!other->m_left || other->m_left->m_color == SS_RBT_BLACK) &&
				(!other->m_right || other->m_right->m_color == SS_RBT_BLACK))
			{
				// Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的  
				other->m_color = SS_RBT_RED;
				node = parent;
				parent = node->m_parent;
			}
			else
			{
				if (!other->m_left || other->m_left->m_color == SS_RBT_BLACK)
				{
					// Case 3: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。  
					other->m_right->m_color = SS_RBT_BLACK;
					other->m_color = SS_RBT_RED;
					LeftRotate(other);
					other = parent->m_left;
				}
				// Case 4: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。
				other->m_color = parent->m_color;
				parent->m_color = SS_RBT_BLACK;
				other->m_left->m_color = SS_RBT_BLACK;
				RightRotate(parent);
				node = m_root;
				break;
			}
		}
	}
	if (node)
		node->m_color = SS_RBT_BLACK;
}

template<typename T_KEY,typename T_POOL>
inline bool ssRBTree<T_KEY,T_POOL>::Check()
{
	SS_ASSERT(m_root->m_parent == 0);
	SS_ASSERT(m_root->m_color);
	if (m_root->m_parent != 0)
		return false;
	if (!m_root->m_color)
		return false;
	if (!CheckRed(m_root))
		return false;
	if (!CheckBlack(m_root,true))
		return false;
	return true;
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::Begin()
{
	m_stack.Empty();
	m_iterator = m_root;
	return Next();
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::End()
{
	return 0;
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::Next()
{
	ssRBTreeNode<T_KEY>* ret = 0;
	while (m_iterator != 0 || m_stack.Count())
	{
		if (m_iterator != 0)
		{
			m_stack.Push(m_iterator);
			m_iterator = m_iterator->Left();
		}
		else
		{
			m_iterator = ret = *m_stack.Pop();
			m_iterator = m_iterator->Right();
			break;
		}
	}
	return ret;
}

template<typename T_KEY, typename T_POOL>
inline void ssRBTree<T_KEY, T_POOL>::CopyFrom(const ssRBTree & t)
{
	RecursionCopy(t.m_root);
}

template<typename T_KEY,typename T_POOL>
inline bool ssRBTree<T_KEY,T_POOL>::CheckRed(ssRBTreeNode<T_KEY>* itr)
{
	if (!itr)
	{
		return true;
	}
    if (itr->Left()!=0)
	{
		if (!CheckRed(itr->Left()))
		{
			return false;
		}
	}
	if (itr->Right()!=0)
	{
		if (!CheckRed(itr->Right()))
		{
			return false;
		}
	}
	if (itr->ColorIsBlack() == SS_RBT_RED)
	{
		if (itr->Left() && itr->Left()->ColorIsBlack() == SS_RBT_RED)
		{
			SS_ASSERT(0);
			return false;
		}
		if (itr->Right() && itr->Right()->ColorIsBlack() == SS_RBT_RED)
		{
			SS_ASSERT(0);
			return false;
		}
		if (itr->Parent() && itr->Parent()->ColorIsBlack() == SS_RBT_RED)
		{
			SS_ASSERT(0);
			return false;
		}
	}
	return true;
}

template<typename T_KEY, typename T_POOL>
inline bool ssRBTree<T_KEY, T_POOL>::CheckBlack(ssRBTreeNode<T_KEY>* itr, bool root)
{
	if (root)
	{
		m_checkBlackCount = 0;
	}
	if (!itr)
	{
		return true;
	}
	if (itr->Left() != 0)
	{
		if (!CheckBlack(itr->Left(),false))
		{
			return false;
		}
	}
	if (itr->Right() != 0)
	{
		if (!CheckBlack(itr->Right(),false))
		{
			return false;
		}
	}
	if (itr->Left()==0 && itr->Right()==0)
	{
		ulint count = itr->ColorIsBlack() ? 1 : 0;
		ssRBTreeNode<T_KEY>* p = itr->Parent();
		while (p)
		{
			if (p->ColorIsBlack())
			{
				++count;
			}
			p = p->Parent();
		}
		if (m_checkBlackCount == 0)
		{
			m_checkBlackCount = count;
		}
		if (count != m_checkBlackCount)
		{
			SS_ASSERT(0);
			return false;
		}
	}

	return true;
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::GetMin()
{
	ssRBTreeNode<T_KEY>* itr = m_root;
	ssRBTreeNode<T_KEY>* index = 0;
	while(itr != 0)
	{
		index = itr;
		if (itr->Left() != 0)
		{
			itr = itr->Left();
		}
		else
		{
			break;
		}
	}
	return index;
}

template<typename T_KEY,typename T_POOL>
inline ssRBTreeNode<T_KEY>* ssRBTree<T_KEY,T_POOL>::GetMax()
{
	ssRBTreeNode<T_KEY>* itr = m_root;
	ssRBTreeNode<T_KEY>* index = 0;
	while(itr != 0)
	{
		index = itr;
		if (itr->Right() != 0)
		{
			itr = itr->Right();
		}
		else
		{
			break;
		}
	}
	return index;
}

template<typename T_KEY,typename T_POOL>
inline void ssRBTree<T_KEY,T_POOL>::Empty()
{
	RecursionEmpty(m_root);
	m_root = 0;
	m_count = 0;
	m_iterator = 0;
	m_stack.Empty();
	m_objPool.Recovery();
}

template<typename T_KEY,typename T_POOL>
void StarSeeker::ssRBTree<T_KEY,T_POOL>::RecursionEmpty(ssRBTreeNode<T_KEY>* itr)
{
	if (itr)
	{
		RecursionEmpty(itr->m_left);
		RecursionEmpty(itr->m_right);
		itr->m_pOnwer = 0;
		m_objPool.Free(itr);
	}
}

template<typename T_KEY, typename T_POOL>
void StarSeeker::ssRBTree<T_KEY, T_POOL>::RecursionCopy(const ssRBTreeNode<T_KEY>* itr)
{
	if (itr)
	{
		Insert(itr->GetKey());
		RecursionCopy(itr->m_left);
		RecursionCopy(itr->m_right);
	}
}

}
#endif
